feat: add hyp migrate (Slurm cluster migration) capability #438
Open
allela-roy wants to merge 1 commit into
Open
feat: add hyp migrate (Slurm cluster migration) capability #438allela-roy wants to merge 1 commit into
allela-roy wants to merge 1 commit into
Conversation
Adds the 'hyp migrate' command for migrating a HyperPod Slurm cluster to a new Availability Zone or Region. Subcommands: snapshot, converge, validate, plan, provision. Bundles an Ansible project (identity + Slurm state capture/ apply) and adds a 'migrate' extra (ansible-core>=2.16).
allela-roy
requested a deployment
to
manual-approval
July 23, 2026 15:28 — with
GitHub Actions
Waiting
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds the
hyp migratecommand for migrating a HyperPod Slurm cluster to a new Availability Zone or Region. Subcommands:snapshot,converge,validate,plan,provision. Bundles an Ansible project (identity + Slurm state capture/apply) and adds amigrateextra (ansible-core>=2.16).What's changing and why?
SMHP Slurm customers provision P* capacity through training plans purchased for a fixed duration. When a plan expires they must buy a new one, and due to capacity constraints the new plan is frequently in a different AZ or Region. Several cluster resources are AZ/Region-scoped and cannot be moved: FSx for Lustre (
/fsx) and FSx for OpenZFS (/home) are single-AZ, the S3 lifecycle-script (LCS) bucket is Region-scoped, and the cluster, its VPC/subnet/security group, and the capacity itself are AZ/Region-scoped.Unless customers manage their clusters with Terraform (not the case ~90% of the time), they must rebuild from scratch with a manual, error-prone runbook. The most common misconfiguration is identity: FSx restores files by numeric owner, so if the new cluster assigns different UIDs/GIDs, every file — including each user's
~/.ssh— is owned by the wrong account and users cannot log in. The rebuild burns significant time and wastes paid training-plan hours while the new plan's clock is already running.This change adds a staged, resumable
hyp migratecommand group that treats a cluster's configuration as a portable artifact and makes the rebuild fast, repeatable, and correct-by-construction:~/.sshpermissions, and reappliesslurm.conf/gres.conf/ accounting./fsx,/home) — the tool creates new FSx for Lustre and OpenZFS filesystems (matching source specs), restores/homefrom a native OpenZFS backup, and carries login-critical~/.sshdata in the manifest. Bulk/fsxdatasets remain the customer's responsibility via their existing Data Repository Associations (DRA).A fail-closed validation gate asserts that target UID/GID and
~/.sshpermissions match the source before cutover, converting "files copied" into "users can actually log in." The tool is shape- and identity-aware (auto-detects fsx-only vs fsx+OpenZFS, and localshared_users.txtvs SSSD/AD) and training-plan aware (TrainingPlanArnbinds per instance group, so the worker group draws from the plan while controller/login run on-demand, with the subnet pinned to the plan's reserved AZ).The
snapshotcapability also doubles as a cost-saving mechanism: capture full cluster state, tear the cluster down to stop compute charges, and recreate it later (same identity, Slurm config, and mounts) when needed.How the Ansible layer reaches nodes
Cluster-internal state (users, groups,
~/.ssh,slurm.conf,gres.conf, accounting) only exists on the live controller. HyperPod nodes aren't standard SSM managed instances, but are reachable through SSM Session Manager via thesagemaker-cluster:<clusterId>_<group>-<instanceId>target. The CLI drives Ansible over theamazon.aws.aws_ssmconnection plugin against a runtime-generated inventory built fromlist-cluster-nodes, giving full idempotent Ansible semantics instead of fragile shell-scraping. Only the controller/login group is targeted; compute nodes self-heal from lifecycle scripts.Before/After UX
Before: No first-class migration path. Customers hand-run a multi-step runbook (recreate filesystems, copy data preserving POSIX ownership, re-upload LCS, recreate the cluster, re-apply Slurm config, clean up). Identity drift silently breaks logins, and reserved training-plan hours are wasted during the rebuild.
After: A five-subcommand workflow inside the existing HyperPod CLI:
Install the optional dependency set with
pip install "sagemaker-hyperpod[migrate]"(addsansible-core); also requires the AWS Session Manager plugin and the collectionsamazon.aws community.aws ansible.posix.How was this change tested?
Validated end-to-end in two scenarios of increasing difficulty. The SSM transport was confirmed in every Region tested (us-east-1, us-east-2, us-west-1): Ansible drives HyperPod controllers over the
aws_ssmconnection plugin on the first try despite HyperPod's login-shell wrapper.Test A — cross-AZ, dual filesystem (us-east-1). Deployed a Slurm cluster with FSx Lustre + OpenZFS in
use1-az1, then migrated touse1-az2. Both healthy;/fsx(Lustre) and/home(OpenZFS / NFS 4.2) mounted; Slurm operational.Test B — cross-Region training-plan migration (us-east-2 → us-west-1). The original intended use case: a Slurm cluster backed by a purchased training plan (reserved P* capacity), migrated to a different Region where a new plan provides capacity.
us-east-2a(use2-az1).us-west-1a(usw1-az3)./fsx) + FSx OpenZFS (/home).Result — cross-Region training-plan migration SUCCEEDED:
idlein Slurm partitiondev./fsx(Lustre) and/home(OpenZFS, NFS 4.2) both mounted.snapshotcaptured identity + Slurm state (shape=fsx_openzfs,identity_model=local, 3 users) to the source-Region S3.aws s3 cp --source-region us-east-2 --region us-west-1).convergereproduced identity exactly (researcher1=5001, researcher2=5002, fsx-users=5000, ubuntu=1000) andvalidatePASSED (after a validator logic fix, below).Are unit tests added?
Yes.
test/unit_tests/cli/test_migrate.py(11 tests) covers themigratecommand group and its subcommands (snapshot,converge,validate,plan,provision) — group wiring, expected subcommand set, and argument handling. All 11 pass locally against currentmain.Are integration tests added?
No automated integration tests are added to CI (they require live HyperPod clusters, training-plan capacity, and multi-Region FSx). The workflow was validated manually end-to-end (Tests A and B above), including the cross-Region training-plan migration on ml.p5.48xlarge. Local run results are summarized in the testing section.
Reviewer Guidelines